# ComfyUI-WanActivationEditor

print("[WanActivationEditor] Loading __init__.py")

from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
from .vector_ops import VECTOR_NODE_CLASS_MAPPINGS, VECTOR_NODE_DISPLAY_NAME_MAPPINGS
from .alternative_nodes import ALTERNATIVE_NODE_CLASS_MAPPINGS, ALTERNATIVE_NODE_DISPLAY_NAME_MAPPINGS
from .block_strength_nodes import NODE_CLASS_MAPPINGS as STRENGTH_NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as STRENGTH_NODE_DISPLAY_NAME_MAPPINGS
from .test_injection_node import NODE_CLASS_MAPPINGS as TEST_NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as TEST_NODE_DISPLAY_NAME_MAPPINGS
from .embedding_amplifier import NODE_CLASS_MAPPINGS as AMPLIFIER_NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as AMPLIFIER_NODE_DISPLAY_NAME_MAPPINGS
from .latent_encoder import NODE_CLASS_MAPPINGS as LATENT_NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as LATENT_NODE_DISPLAY_NAME_MAPPINGS
from .projection_booster import NODE_CLASS_MAPPINGS as BOOSTER_NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as BOOSTER_NODE_DISPLAY_NAME_MAPPINGS

# Import SAE nodes if available
try:
    from .sae_nodes import NODE_CLASS_MAPPINGS as SAE_NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as SAE_NODE_DISPLAY_NAME_MAPPINGS
    print("[WanActivationEditor] SAE nodes loaded successfully")
except ImportError as e:
    print(f"[WanActivationEditor] SAE nodes not loaded (optional): {e}")
    SAE_NODE_CLASS_MAPPINGS = {}
    SAE_NODE_DISPLAY_NAME_MAPPINGS = {}

# Merge all node mappings
ALL_NODE_CLASS_MAPPINGS = {
    **NODE_CLASS_MAPPINGS,
    **VECTOR_NODE_CLASS_MAPPINGS,
    **ALTERNATIVE_NODE_CLASS_MAPPINGS,
    **STRENGTH_NODE_CLASS_MAPPINGS,
    **TEST_NODE_CLASS_MAPPINGS,
    **AMPLIFIER_NODE_CLASS_MAPPINGS,
    **LATENT_NODE_CLASS_MAPPINGS,
    **BOOSTER_NODE_CLASS_MAPPINGS,
    **SAE_NODE_CLASS_MAPPINGS,
}
ALL_NODE_DISPLAY_NAME_MAPPINGS = {
    **NODE_DISPLAY_NAME_MAPPINGS,
    **VECTOR_NODE_DISPLAY_NAME_MAPPINGS,
    **ALTERNATIVE_NODE_DISPLAY_NAME_MAPPINGS,
    **STRENGTH_NODE_DISPLAY_NAME_MAPPINGS,
    **TEST_NODE_DISPLAY_NAME_MAPPINGS,
    **AMPLIFIER_NODE_DISPLAY_NAME_MAPPINGS,
    **LATENT_NODE_DISPLAY_NAME_MAPPINGS,
    **BOOSTER_NODE_DISPLAY_NAME_MAPPINGS,
    **SAE_NODE_DISPLAY_NAME_MAPPINGS,
}

# Export for ComfyUI
NODE_CLASS_MAPPINGS = ALL_NODE_CLASS_MAPPINGS
NODE_DISPLAY_NAME_MAPPINGS = ALL_NODE_DISPLAY_NAME_MAPPINGS

WEB_DIRECTORY = "./js"

__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS', WEB_DIRECTORY]
